package com.lwd.util.collection;

import java.util.*;
import java.util.function.Predicate;
import java.util.function.Supplier;

/**
 * @author Brozen
 * @date 2019/6/11 2:40 PM
 */
public class EnhancedCollectionUtils {


    public static <E, K> Map<K, E> map(Collection<E> collection, Function<E, K> keyGenerator) {
        return map(collection, keyGenerator, new Function<E, E>() {
            @Override
            public E apply(E input) {
                return input;
            }
        });
    }

    public static <E, K, V> Map<K, V> map(Collection<E> collection, Function<E, K> keyGenerator,
                                                               Function<E, V> valueGenerator) {
        return map(collection, keyGenerator, valueGenerator, new Supplier<Map<K, V>>() {
            @Override
            public Map<K, V> get() {
                return new HashMap<>();
            }
        });
    }

    public static <E, K, V, M extends Map<K, V>> Map<K, V> map(Collection<E> collection, Function<E, K> keyGenerator,
                                                               Function<E, V> valueGenerator, Supplier<M> collectionSupplier) {
        Assert.notNull(collection);
        Assert.notNull(keyGenerator);
        Assert.notNull(valueGenerator);
        Assert.notNull(collectionSupplier);

        M resultMap = collectionSupplier.get();
        for (E e : collection) {
            K key = keyGenerator.apply(e);
            V value = valueGenerator.apply(e);
            resultMap.put(key, value);
        }

        return resultMap;
    }

    public static <T, E> Map<T, Collection<E>> group(Collection<E> collection, Function<E, T> keyGenerator) {
        return group(collection, keyGenerator, new Supplier<Collection<E>>() {
            @Override
            public Collection<E> get() {
                return new LinkedList<>();
            }
        });
    }

    public static <T, E, C extends Collection<E>> Map<T, C> group(Collection<E> collection, Function<E, T> keyGenerator, Supplier<C> collectionSupplier) {
        Assert.notNull(keyGenerator);
        Assert.notNull(collectionSupplier);

        Map<T, C> groupedResult = new HashMap<>();
        for (E e : collection) {
            T key = keyGenerator.apply(e);
            if (!groupedResult.containsKey(key)) {
                C coll = collectionSupplier.get();
                groupedResult.put(key, coll);
            }

            groupedResult.get(key).add(e);
        }

        return groupedResult;
    }

    public static <I, O> Collection<O> transform(final Collection<I> input, Function<I, O> transformer) {
        return transform(input, transformer, new Supplier<Collection<O>>() {
            @Override
            public Collection<O> get() {
                return new ArrayList<>(input.size());
            }
        });
    }

    public static <I, O, C extends Collection<O>> C transform(Collection<I> input, Function<I, O> transformer, Supplier<C> collectionSupplier) {
        Assert.notNull(input);
        Assert.notNull(transformer);
        Assert.notNull(collectionSupplier);

        C result = collectionSupplier.get();
        for (I i : input) {
            result.add(transformer.apply(i));
        }
        return result;
    }


    public static <E, C extends Collection<E>> C filter(C list, Predicate<E> predicate) {
        Assert.notNull(list);
        Assert.notNull(predicate);

        Iterator<E> iterator = list.iterator();
        while (iterator.hasNext()) {
            if (predicate.apply(iterator.next())) {
                iterator.remove();
            }
        }
        return list;
    }

    public static <E, C extends Collection<E>> C filter(Collection<E> list, Predicate<E> predicate, Supplier<C> collectionSupplier) {
        Assert.notNull(list);
        Assert.notNull(predicate);
        Assert.notNull(collectionSupplier);

        C result = collectionSupplier.get();
        for (E ele : list) {
            if (!predicate.apply(ele)) {
                result.add(ele);
            }
        }
        return result;
    }

    public static <E, R> R reduce(Collection<E> list, Reducer<E, R> reducer) {
        return reduce(list, reducer, null);
    }

    public static <E, R> R reduce(Collection<E> list, Reducer<E, R> reducer, R identity) {
        Assert.notNull(list);
        Assert.notNull(reducer);

        R result = identity;
        for (E e : list) {
            result = reducer.apply(e, result);
        }

        return result;
    }
}
